'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: retrieval_service.py
* @Time: 2025/9/28
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass
from uuid import UUID

from flask import Flask
from langchain.retrievers import EnsembleRetriever
from langchain_core.documents import Document as LCDocument
from injector import inject
from langchain_core.tools import BaseTool, tool
from pydantic import BaseModel, Field
from sqlalchemy import update
from pkg.sqlachemy.sqlalchemy import SQLAlchemy
from internal.entity.dataset_entity import RetrievalStrategy, RetrievalSource
from internal.model import Dataset, DatasetQuery, Segment
from internal.exception.exception import NotFoundException
from .jieba_service import JiebaService
from .vector_database_service import VectorDatabaseService
from .base_service import BaseService
from ..lib.helper import combine_documents


@inject
@dataclass
class RetrievalService(BaseService):
    """检索服务"""
    db:SQLAlchemy
    jieba_service:JiebaService
    vector_database_service:VectorDatabaseService

    def search_in_datasets(self, dataset_ids:list[UUID],
                           query:str,
                           account_id:UUID,
                           retrieval_strategy:str = RetrievalStrategy.SEMANTIC,
                           k:int=4,
                           score:float = 0,
                           retrieval_source:str = RetrievalSource.HIT_TESTING)->list[LCDocument]:
        #1、提取知识库列表并校验授权认证
        datasets = self.db.session.query(Dataset).filter(
            Dataset.id.in_(dataset_ids),
            Dataset.account_id == account_id
        ).all()
        if datasets is None or len(datasets) == 0:
            raise NotFoundException("当前 无知识库可以检索！！")
        dataset_ids = [dataset.id for dataset in datasets]

        #2、构建不同种类的检索器
        from internal.core.retrievers import SemanticRetriever,FullTextRetriever
        semantic_retriever = SemanticRetriever(
            dataset_ids = dataset_ids,
            vector_store=self.vector_database_service.vector_store,
            search_kwargs={
                "k":k,
                "score_threshold":score
            }
        )

        full_text_retriever = FullTextRetriever(
            db = self.db,
            dataset_ids=dataset_ids,
            jieba_service=self.jieba_service,
            search_kwargs={
                "k":k,
            }
        )

        hybrid_retriever = EnsembleRetriever(
            retrievers=[semantic_retriever, full_text_retriever],
            weights=[0.5, 0.5]
        )

        #3、根据不同的检索策略执行不同的检索
        if retrieval_strategy == RetrievalStrategy.SEMANTIC:
            lc_documents = semantic_retriever.invoke(query)[:k]
        elif retrieval_strategy == RetrievalStrategy.FULL_TEXT:
            lc_documents = full_text_retriever.invoke(query)[:k]
        else:
            lc_documents = hybrid_retriever.invoke(query)[:k]

        #4、将检索 记录添加到数据库中
        unique_dataset_ids = list(set(str(lc_document.metadata["dataset_id"]) for lc_document in lc_documents))
        for dataset_id in unique_dataset_ids:
            self.create(DatasetQuery,
                        dataset_id = dataset_id,
                        query=query,
                        source=retrieval_source,
                        source_app_id = None,#//todo 后续完善功能后再调整
                        created_by=account_id
            )

        #5、批量更新片段命中的次数
        with self.db.auto_commit():
            stmt = (
                update(Segment)
                .where(Segment.id.in_([lc_document.metadata["segment_id"]for lc_document in lc_documents]))
                .values(hit_count = Segment.hit_count + 1)
            )
            self.db.session.execute(stmt)

        return lc_documents



    def create_langchain_tool_from_serch(self,
                                         flask_app:Flask,
                                         dataset_ids:list[UUID],
                                         account_id:UUID,
                                         retrieval_strategy:str = RetrievalStrategy.SEMANTIC,
                                         k:int=4,
                                         score:float = 0,
                                         retrieval_source:str = RetrievalSource.HIT_TESTING,
                                         )->BaseTool:
        """根据传递的参数构建一个Langchain知识库搜索工具"""
        class DatasetRetrievalInput(BaseModel):
            """知识库检索工具输入结构"""
            query:str = Field(description="知识库搜索query语句,类型为字符串！")

        @tool
        def dataset_retrieval(query:str)->str:
            """如果需要搜索扩展的知识库内容，当你觉得用户的问题超过你的知识范围时， 可以尝试调用该工具，输入为搜索的query语句，返回数据为检索内容字符串"""
            #1、调用search_in_datasets检得到Langchian 文档列表
            with flask_app.app_context():
                documents = self.search_in_datasets(
                        dataset_ids = dataset_ids,
                        query=query,
                        account_id=account_id,
                        retrieval_strategy=retrieval_strategy,
                        k=k,
                        score=score,
                        retrieval_source=retrieval_source,
                )
            #2、将Langchain文档列表转换成字符串后返回
            if len(documents) == 0:
                return "知识库没有检索大对应的内容"
            return combine_documents(documents)

        return dataset_retrieval



